#!/usr/bin/python
###############################################################################
#                                                                             #
#    Copyright 2007 Danil Dotsenko                                       #
#                                                                             #
#    This program is free software; you can redistribute it and/or modify     #
#    it under the terms of the GNU General Public License as published by     #
#    the Free Software Foundation; either version 2 of the License, or        #
#    (at your option) any later version.                                      #
#                                                                             #
#    This program is distributed in the hope that it will be useful,          #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of           #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the            #
#    GNU General Public License for more details.                             #
#                                                                             #
#    You should have received a copy of the GNU General Public License        #
#    along with this program; if not, write to the Free Software              #
#    Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307  USA #
#                                                                             #
###############################################################################
# Contacts: Danil Dotsenko, dd@accentsolution.com
# version 20070505

import time
from socket import *
import re
from urllib import urlopen

'''
python-based uPnP library.
Allows a variety of actions on the default WANConnection device
Most often-used actions (get-External-IP, add-Port, del-Port, 
list-Ports and test-Specific-Port) are made into functions
for ease of usage.
Note. This is a LAZY library and relies on you giving some input.
To activate the library you need: ClientMachineIP, RoutersInternalIP
Example:
>>> import upnpc
>>> upnpclient.client.activate('192.168.1.104','192.168.1.1') # our Internal IP, gateway IP
This will talk to the gateway (192.168.1.1) and import all needed
pieces from it's control set.
Then just call: (note the optional = default values in [])
>>> upnpclient.client.getExtIP([refresh=False - may return cached value]) for external IP
>>> upnpclient.client.addMapping(external_port [, protocol="TCP", ip=this_machine's, intport=same_as_ext])
>>> upnpclient.client.delMapping(external_port [, protocol="TCP"])
>>> upnpclient.client.listMapping()
>>> upnpclient.client.testPort(external_port [, protocol="TCP"]) returns 200 if assigned, else, available.
'''

class upnpclient:
	def __init__(self):
		#self.WANConnection = None # will be class instance with WAN[IP|PPP]Connection properties
		self.info = 'custom' # this text will appear as "explanation" or "application name" in the list next to the assigned port when you look on router web ui.
		self.extIP = ''
		self.intIP = ''
		class SuperObjectClass:
			def __init__(self, parent=None):
				self.objtype = 'SuperObjectClass'
		self.service = SuperObjectClass()
		self.service.addr = ''
		self.service.port = 80
		self.service.typeid = None # 'urn:schemas-upnp-org:service:WANIPConnection:1' # serviceType tag
		self.service.querrypath = None # '"/WANIPConnection"' # a.k.a "selector" arg for putrequest()

	def activate(self, intip, gatewayip):
		self.intIP = intip # a string, '123.456.778.432'
		self.service.addr = gatewayip # '192.168.1.1'
		text = self._greeting_()
		if text:
			return self._handshake_(text) # which also may return False, which statns for Failure.
		else:
			return False

	def _isolatenext_(self, all, textend, limit):
		'''
		Finds range of the tag, based on start / closing tags.
		Respects nested tags of the same type and skips their closing tag.
		Returns either (starting address, ending address) in the string, or (None, None)
		'''
		# separating the 1st out.
		block = []
		que = 0
		tstart, tend = None, None
		end = False
		while not end:
			try:
				i = all.next() # note, we already took the 1st one 
			except:
				i = None
			if not i:
				end = True # hmm, abrupt end. Likely, no closing tag.
				tend = textend # we just append the rest ot the text as part ot the tag.
			else:
				if i.group()[1] == '/': # i.e. if it's a closing tag, "</blah>"
					que = que - 1
					if que-1 == limit:
						block.append(i.end())
					if que < 1: # this is the end, my lonely end, my friend.
						tend = i.start() # we don't need the tag itself, see.
						end = True
					# print 'minus ', que, i.group()
				else: # ugh! another nested tag of the same type.
					que = que + 1
					if tstart == None:
						tstart = i.end() # we don't need the tag itself, see.
					if que - 2 == limit:
						block.append(i.start())
					# print 'plus ', que, i.group()
		return tstart, tend, block
	
	def getsections(self, tag, text, limit=None):
		'''
		Returns a list of tag sections in XML-style formatted text
		example: upnpc.getsections('as','<as>qqqqqqqq</as>aaaaaaa<as>zzzzzzzzz</as>ccccccc<as>yyyyyy')
		returns ['qqqqqqqq','zzzzzzzzz']
		Note, if openning tag has no closing tag, the last element may thrown away. Occupational hazards.
		It may be just the <tag>, or the tag with the remainder of the text contents.
		- tag - is the tag we are looking for
		- text - the text where we are looking for the doc.
		- limit -  how many NESTED levels of the same tag we allow to fall into the
		resulting text snippet. Usefull to filter out nested tags of same type.
		Defaults to no limit. 0 does not allow nested tags. 1 allows 1 nested level. etc.
		'''
		sch = re.compile('</'+tag+'>|<'+tag+'[^>]*>') # the last one allows in-tag options in openning tag
		all = sch.finditer(text)
		snippets = []
		
		textend = text.__len__()
		st, en, block = self._isolatenext_(all, textend, limit) # note, this can return (0, 4) = so, dont just test for bool
		while st!=None and en!=None:
			snippets.append((st, en, block))
			st, en, block = self._isolatenext_(all, textend, limit)
		# finally, assembling the answers
		ret = []
		for i in snippets:
			blklen = i[2].__len__()
			if blklen: # i.e. if we have blocking areas
				# ugh, building the snippet from parts.
				if blklen != round(blklen/2)*2: # there is no built-in  method for odd/even i could find.
					# if we are here, the blackout ranges are odd. Removing last item.
					del i[2][blklen-1] 
			rnge = [i[0]]
			rnge.extend(i[2])
			rnge.append(i[1])
			# as a result, we have a sandwitch: 
			# [tag_start, blk1_st, blk1_end, blk2_st, blk2_end, tag_end]
			# or simply [tag_start, tag_end] if there was no blocking
			txt = ''
			# taking a pair of limits, and adding section that
			# corresponds to that pair to output
			for i in range(int(rnge.__len__()/2)):
				txt = txt+text[rnge[i*2]:rnge[(i*2)+1]]
			if txt: # sometimes, like if there was only one openning tag, we may have answers like this: "" - no use to us.
				ret.append(txt)
		return ret

	def _greeting_(self):
		'''
		Getting the info out of the gateway device. 
		Note, we expect to be told what the gateway ip is,
		otherwise we will not know which of the upnp devices to talk to.
		'''
		# getting "greetings" ouf the gateway
		s = socket(AF_INET, SOCK_DGRAM) # this means we are in UDP mode
		s.setsockopt(SOL_SOCKET, SO_BROADCAST, 1) # this means we can spew out broadcasts
		s.bind(('', 5000)) # i think this should be port 5000... not sure, but 0 works.		
		text = 'M-SEARCH * HTTP/1.1\r\nHost:239.255.255.250:1900\r\nST:upnp:rootdevice\r\nMan:"ssdp:discover"\r\nMX:3\r\n\r\n'
		s.setblocking(0)
		s.sendto(text, ('239.255.255.250', 1900))
		t = time.time()+2 # we give all devices 2 seconds to respond.
		greet = None
		while time.time() < t and not greet:
			time.sleep(0.25) # 1/4 of a second delay. No reason to rush up 100 of CPU here.
			try:
				msg, addr = s.recvfrom(4096) #, MSG_DONTWAIT) # this way we don't block if there is no output, we raise error
				if addr[0]==self.service.addr:
					greet = msg
			except:
				pass
		s.close()
		# now that we have a greeting, let's get the full introduction.
		# grrr... we hates \r\n at the end of greetingses.
		# greet = greet.replace('\r\n','')
		msg = ''
		if greet:
			# parse the introduction request URL out of the greeting.
			# it will always be after Location:
			try:
				url = re.search('(?<=Location:).+', greet).group().strip()
			except:
				url = None
			if url:
				d = urlopen(url)
				msg = d.read()
				d.close()
				del d
		msg = msg.replace('\r\n','')
		return msg

	def _handshake_(self, devinfo):
		'''
		devinfo  - xml with gateway's greeting
		Here we pick it apart, testing the gateway for it's functionality.
		And hand it over to interested party.
		Items of interest in the XML document:
		root
		  |_ InternetGatewayDevice (device section)
				|_ WANDevice (device section)
					|_ WANConnectionDevice (device section)
						|_ WAN[PPP|IP]ConnectionService (service section)
						    for PPPoE / dialup - WANPPPConnectService
							for normal DSL/Cable connections - WANIPConnectService
		URLBase - at root level, unique, is the base url for all http XML requests
		'''
		
		# 1. Lets see it it's actually an InternetGatewayDevice
		devices = self.getsections('deviceType',devinfo)
		if ''.join(devices).find('WANConnectionDevice') == -1:
			# we dont see what we need
			devinfo = ''
		del devices
		
		def contains(text, snippets):
			result = ''
			for i in snippets:
				if i.find(text) != -1:
					result = i
					break
			return result

		# we are looking for a device section that doesn't contain "WANDevice"
		# but contains "WANConnectionDevice"
		snippet = devinfo
		while snippet:
			devinfo = self.getsections('device', snippet)
			snippet = contains('WANDevice', devinfo)
			# whichever the 1st devinfo will not contain "WANDevice"
			# is the one that will have our "WANConnectionDevice"
			# but, we will test anyway, using residual devinfo
		snippet = contains('WANConnection', devinfo)
		
		service = ''
		if snippet:
			# Now, we are in the right place.
			# squeeze out the WAN[IP|PPP]Connection service entries.
			services = self.getsections('service', snippet)
			if services.__len__() == 1:
				# nice! this is simple. It's possibly the only one active.
				service = services[0]
			elif wanconn.__len__() > 1:
				# errrr, demn it. I'll take the easy route here.
				# later I'll test for which one is actually active.
				service = services[0]

		# 2. Now, we take the service apart for scrap.
		# We need the following pieces:
		# a. how to find out the outside IP
		# b. how to request add to forward list
		# c. how to request delete from forward list
		# d. how to find out the list of forwarded ports.
		# e. how to find out if our desired port is occupied.
		# that's all we care about.
		# this is an example of what we are looking at:
		# <serviceType>urn:schemas-upnp-org:service:WANIPConnection:1</serviceType>
		# <serviceId>urn:upnp-org:serviceId:WANIPConnection</serviceId>
		# <controlURL>http://192.168.1.1:2468//WANIPConnection</controlURL>
		# <eventSubURL>/WANIPConnection</eventSubURL>
		# <SCPDURL>/WANIPCn.xml</SCPDURL>

		if service:
			try:
				self.service.typeid = self.getsections('serviceType',service)[0]
				controlURL = self.getsections('controlURL',service)[0]
			except:
				controlURL = ''
				# leave typeid None - we'll check for it all the time and refresh again.
			if controlURL:
				# examples:
				# (netgear) '/uuid:000625d8-095f-0006-25d8-095f0232aa18/WANPPPConnection:1'
				# (linksys) 'http://192.168.1.1:2468//WANIPConnection'
				# not a friendly bunch. :(
				# 1. lets test for http:// if we have it, the port is different.
				if controlURL[0:7] == 'http://':
					controlURL = controlURL[7:]
					try:
						address, path = controlURL.split('/',1)
					except:
						address, path = '','/'
					try:
						trash, port = address.split(':')
						port = int(port)
					except:
						port = 80
					# by spliting at / we removed it. need it back.
					if path[0] != '/':
						path = '/'+path
					self.service.port = port
					self.service.querrypath = path
				else:
					self.service.querrypath = controlURL
		if self.service.querrypath:
			return True
		else:
			return False

	def action(self, action='GetExternalIPAddress', args={}, service=None):
		'''
		Action is just a text name of the action
		Args is a dict type of arguments.
		Service is our custom-assembled class instance of one of the upnp services.
			these are created when the gateway is greeted and embraced.
			By default it is either WANIPConnection or WANPPPConnection
			whichever was detected as active.
		We return error codes.
		0 - no answer was returned in 1.5 seconds.
		200 - everything is OK.
		all other codes are passed directly from the uPnP reply message.
		See associated upnp keywords for error codes that relate to them
		'''

		if not self.intIP:
			return '', 0

		body1 = '<?xml version=\"1.0\" encoding=\"utf-8\"?>\n<s:Envelope xmlns:s="http://schemas.xmlsoap.org/soap/envelope/"\ns:encodingStyle="http://schemas.xmlsoap.org/soap/encoding/">\n<s:Body>\n<u:%s xmlns:u="%s">\n' % (action, service.typeid)
		body2 = '</u:%s>\n</s:Body>\n</s:Envelope>\n\n' % action
		#example args = {'NewRemoteHost':'', 'NewProtocol':'TCP', 'NewExternalPort':8080}
		arguments = ''
		for arg in args:
			arguments = arguments + ('<%s>%s</%s>\n' % (arg, args[arg], arg))
		body = body1+arguments+body2

		h = '''POST "%s" HTTP/1.1
HOST: %s
CONTENT-TYPE: text/plain; charset="utf-8"
CONTENT-LENGTH: %s
SOAPACTION: "%s#%s"''' % (service.querrypath, (service.addr+str(service.port)), body.__len__(), service.typeid, action)

		s = socket(AF_INET,SOCK_STREAM)
		s.connect((service.addr,service.port))
		s.send(h+body)
		t = time.time()+1
		msg = ''
		while time.time() < t and not msg:
			time.sleep(0.1) # 1/4 of a second delay. No reason to rush up 100 of CPU here.
			try:
				msg = s.recv(4096,MSG_DONTWAIT) # this way we don't block if there is no output, we raise error
			except:
				pass
		s.close()
		
		try:
			h, status, trash = msg.split(' ', 2)
			status = int(status)
		except:
			status = 0

		if status != 200:
			# upnp requires to ask as post first, and 
			# if there is an error, ask with special method
			# note, we are just changing the header
			h = '''M-POST "%s" HTTP/1.1
HOST: %s
CONTENT-TYPE: text/plain; charset="utf-8"
CONTENT-LENGTH: %s
MAN: "http://schemas.xmlsoap.org/soap/envelope/"; ns=01
01-SOAPACTION: "%s#%s"''' % (service.querrypath, (service.addr+str(service.port)), body.__len__(), service.typeid, action)
			s = socket(AF_INET,SOCK_STREAM)
			s.connect((service.addr,service.port))
			s.send(h+body)
			t = time.time()+1
			msg = ''
			while time.time() < t and not msg:
				time.sleep(0.1) # 1/4 of a second delay. No reason to rush up 100 of CPU here.
				try:
					msg = s.recv(4096,MSG_DONTWAIT) # this way we don't block if there is no output, we raise error
				except:
					pass
			s.close()
			try:
				h, status, trash = msg.split(' ', 2)
				status = int(status)
			except:
				status = 0

		body = ''
		if status:
			i = re.search('<\?xml version="[\d\.]+"\?>',msg)
			if i:
				body = msg[i.end():]
			body = body.replace('\r\n','') # I hope this doesn't bite me later.
		#print "body, status are ",body, status
		return body, status

	def getExtIP(self, refresh=False):
		'''
		Provides previously querried from the gateway external IP number.
		If any non-Null argument is passed, refreshes the number.
		'''
		ip = ""
		if self.extIP and not refresh:
			ip = self.extIP
		else:
			# put some checks here that self.WANConnection is ok.
			body, status = self.action(action='GetExternalIPAddress', args={}, service = self.service)
			if body:
				# that's a chug full of XML.
				try:
					self.extIP = self.getsections('NewExternalIPAddress', body)[0]
				except:
					pass
		return self.extIP

	def addMapping(self, extport, ip=None, protocol="TCP", intport=None, enabled = '1'):
		'''
		Error codes:
		200 OK, all is well
		402 Invalid Args
		501 Action Failed
		715 WildCardNotPermittedInSrcIP
		716 WildCardNotPermittedInExtPort
		718 ConflictInMappingEntry
		724 SamePortValuesRequired
		725 OnlyPermanentLeasesSupported
		726 RemoteHostOnlySupportsWildcard
		727 ExternalPortOnlySupportsWildcard
		Watch for 718 in particular and try new port.
		'''
		if not intport:
			# some routers can accept trafic on one external port and send it to other internal port.
			# but most of the time, we use the same port.
			intport = extport
		if not ip:
			ip = self.intIP
		args = {#'NewRemoteHost':'',
			'NewExternalPort':extport,
			'NewProtocol':protocol,
			'NewInternalPort':intport,
			'NewInternalClient':ip,
			'NewEnabled':enabled,
			'NewPortMappingDescription':self.info,
			#'NewLeaseDuration':'',
			}
		body, status = self.action('AddPortMapping', args, service = self.service)
		if status !=200:
			#extract the error out of the body.
			try:
				status = int(self.getsections('errorCode', body)[0])
			except:
				pass
		return status

	def delMapping(self, port, protocol='TCP'):
		'''
		Error codes:
		200 OK, all is well
		402 Invalid Args
		714 NoSuchEntryInArray
		'''
		args = {#'NewRemoteHost':'',
			'NewExternalPort':port,
			'NewProtocol':protocol,}
		#print args
		body, status = self.action('DeletePortMapping', args, service = self.service)
		if status !=200:
			#extract the error out of the body.
			try:
				status = int(self.getsections('errorCode', body)[0])
			except:
				pass
		return status

	def listMapping(self):
		'''
		Produces a list of dict types with the followinf contents:
		[{'info': 'KTorrent', 'intport': 4444, 'extport': 4444, 'protocol': 'UDP', 'enabled': 1, 'intIP': '192.168.1.104'},{another dyct type}]
		'''
		status = 200
		counter = 0
		args = {'NewPortMappingIndex':''}
		results = []
		while status == 200 and counter < 64: # 64 is arbitrary. My guess of max supported uPnP assignment slots for a router.
			args.update({'NewPortMappingIndex':counter})
			body, status = self.action('GetGenericPortMappingEntry', args, service = self.service)
			counter += 1
			# need: extport, type, IP, intport, enabled.
			if status == 200:
				try:
					a = {}
					a['extport'] = int(self.getsections('NewExternalPort',body)[0])
					a['protocol'] = self.getsections('NewProtocol',body)[0]
					a['intport'] = int(self.getsections('NewInternalPort',body)[0])
					a['intIP'] = self.getsections('NewInternalClient',body)[0]
					a['enabled'] = int(self.getsections('NewEnabled',body)[0])
					a['info'] = self.getsections('NewPortMappingDescription',body)[0]
					results.append(a.copy())
				except:
					pass
		return results
	
	def testPort(self, port, protocol='TCP'):
		'''
		Checks for assignment on the specific port.
		Returns a tuple of
		- Dict object with settings for the port.
		- status code. 200 = OK. anything else.. handle it.
		'''
		args = {#'NewRemoteHost':'',
			'NewExternalPort':port,
			'NewProtocol':protocol}
		body, status = self.action('GetSpecificPortMappingEntry', args, service = self.service)
		a = {}
		if status !=200:
			#extract the error out of the body.
			try:
				status = int(self.getsections('errorCode', body)[0])
			except:
				pass
		else:
			#a['extport'] = int(self.getsections('NewExternalPort',body)[0])
			#a['protocol'] = self.getsections('NewProtocol',body)[0]
			a['intport'] = int(self.getsections('NewInternalPort',body)[0])
			a['intIP'] = self.getsections('NewInternalClient',body)[0]
			a['enabled'] = int(self.getsections('NewEnabled',body)[0])
			a['info'] = self.getsections('NewPortMappingDescription',body)[0]
		return a, status

client = upnpclient()